import requests
import json

# Ollama API 地址
OLLAMA_API_URL = "http://localhost:11434"
MODEL_NAME = "deepseek-r1:1.5b"

# 系统提示
SYSTEM = """
你是一个严谨的AI助手，必须用中文回答问题。
你的回答必须详细，准确，并且遵循以下格式：
User: [用户输入]
Assistant: [你的回答]
如果用户的问题不明确，请要求用户提供更多信息。
"""

# 模板
def build_prompt(system_prompt, user_input):
    return f"""<|system|>
{system_prompt}<|end|>
<|user|>
User: {user_input}<|end|>
<|assistant|>
Assistant:"""

# 用户输入
# user_input = "你好，DeepSeek！"
user_input = "你知道水稻有哪些害虫吗？"

# 组织请求数据
data = {
    "model": MODEL_NAME,
    "prompt": build_prompt(SYSTEM, user_input),  # 使用模板构建完整的提示
}

try:
    # 发送 POST 请求
    print("正在发送请求...")
    response = requests.post(f"{OLLAMA_API_URL}/api/generate", json=data, timeout=30, stream=True)
    print("请求完成，状态码：", response.status_code)

    # 打印 API 响应
    if response.status_code == 200:
        # print("模型回复：", response.text)  # 直接打印原始响应

        full_response = ""

        for line in response.iter_lines(): 
            if line:
                # 解析 JSON
                chunk = line.decode("utf-8")  # 将字节流转换为字符串
                json_data = json.loads(chunk)  # 将字符串转换为字典

                if "response" in json_data:
                    full_response += json_data["response"]
                if json_data.get("done", False):
                    break  # 如果 done 为 True，结束循环

        print("完整回复：", full_response)


    else:
        print(f"请求失败，状态码：{response.status_code}")
        print("API 返回：", response.text)
except requests.exceptions.Timeout:
    print("请求超时，请检查网络或 API 服务是否正常。")
except requests.exceptions.RequestException as e:
    print(f"请求发生错误：{e}")

